import os
import pandas as pd

CASIA_SAVE_FOLDER = '../data/cssia_fasd/'


def generate_csv(root_dir, mode):
    if 'train' in mode:
        videos_path = os.path.join(root_dir, 'train_release_crops')
        # videos_path = os.path.join(root_dir, 'Train_files_MSR')
    else:
        videos_path = os.path.join(root_dir, 'test_release_crops')
        # videos_path = os.path.join(root_dir, 'Test_files_MSR')
    fold_data = []
    for video_name in os.listdir(videos_path):
        video_name_split = video_name.split('_')
        video_path = os.path.join(videos_path, video_name)
        for image in os.listdir(video_path):
            if len(video_name_split) == 2:
                if video_name_split[1] == '1' or video_name_split[1] == '2':
                    label = 0
                    fold_data.append([video_name, image, label])
                else:
                    label = 1
                    fold_data.append([video_name, image, label])
            elif len(video_name_split) == 3:
                if video_name_split[2] == '1':
                    label = 0
                    fold_data.append([video_name, image, label])
                else:
                    label = 1
                    fold_data.append([video_name, image, label])
    columns = ["video", "file", "label"]
    pd.DataFrame(fold_data, columns=columns).to_csv(f'{CASIA_SAVE_FOLDER}data_{mode}.csv', index=False)
    # pd.DataFrame(fold_data, columns=columns).to_csv('../data/data_{}_MSR.csv'.format(mode), index=False)


def main():
    os.makedirs(CASIA_SAVE_FOLDER, exist_ok=True)
    # root_dir = '/data/heyan/Datasets/Face/oulu_256'  # RGB images
    # root_dir = '/home/shaohua/data2/Datasets/Oulu_NPU_256'  # MSR images
    root_dir = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/CASIA_FASD/casia'

    generate_csv(root_dir, 'train')
    # generate_csv(root_dir, 'val')
    generate_csv(root_dir, 'test')

    # generate_csv(root_dir, 'train_real')
    # generate_csv(root_dir, 'train_fake')


if __name__ == '__main__':
    main()
